# Copyright (c) 2021 Presto Labs Pte. Ltd.
# Author: xguo

import datetime
import json
import os
import subprocess
import sys

from absl import app, flags

from coin.support.proto_log.logic.util import run_from_og_log_archive
import google.protobuf.json_format as json_format

import coin.proto.coin_order_enums_pb2 as coin_enum
from coin.proto.coin_query_pb2 import (
    AccountBalance,
    AccountPosition,
    ProductFillBundle,
    ProductOrderBundle,
    ConversionBundle,
    SubmitOrderResponse,
    CancelOrderResponse,
    RawMessageProto,
    FundingFeeBundle,
    OrderExecutionProto,
    MultiSubmitResponse,
    MultiCancelResponse,
    ErrorCode,
)
from coin2.service.web3.order_service_pb2 import QueryNftTokenInfoResponseProto
from coin.proto.coin_order_gateway_pb2 import OrderGatewayLog

FLAGS = flags.FLAGS


class OnOpLog(object):
  def __init__(self, filename):
    self.file_obj = open(filename, "w")
    self.last_print_time = 0
    self.num = 0

  def close(self):
    self.file_obj.close()

  def on_log(self, timestamp, log):
    pb = OrderGatewayLog()
    pb.ParseFromString(log)
    source = coin_enum.OrderSubTopicId.Name(pb.raw_msg_id >> (64 - 8))
    if pb.type == OrderGatewayLog.PRIVATE_EXCHANGE_MESSAGE:
      raw_msg = RawMessageProto()
      raw_msg.ParseFromString(pb.private_exchange_message)

      if raw_msg.HasField('raw_bytes'):
        raw_bytes = raw_msg.raw_bytes
        raw_msg.ClearField('raw_bytes')

      if raw_msg.HasField('http_response'):
        res_body = raw_msg.http_response.body
        if raw_msg.http_response.HasField('http_request'):
          req_body = raw_msg.http_response.http_request.body
          raw_msg.http_response.http_request.ClearField('body')
        raw_msg.http_response.ClearField('body')

      if raw_msg.HasField('http_context'):
        res_body = raw_msg.http_context.http_response.body
        req_body = raw_msg.http_context.http_request.body
        raw_msg.http_context.http_request.ClearField('body')
        raw_msg.http_context.http_response.ClearField('body')

      raw_msg_dict = json_format.MessageToDict(raw_msg, preserving_proto_field_name=True)
      raw_msg_dict['type'] = coin_enum.OrderSubTopicId.Name(raw_msg.type)
      msg_type = raw_msg_dict['type']

      if msg_type == 'PARSED_ACCOUNT_BALANCE':
        account_balance = AccountBalance()
        account_balance.ParseFromString(raw_bytes)
        raw_msg_dict['raw_bytes'] = json_format.MessageToDict(
            account_balance, preserving_proto_field_name=True)
      elif msg_type == 'PARSED_ACCOUNT_POSITION':
        account_pos = AccountPosition()
        account_pos.ParseFromString(raw_bytes)
        raw_msg_dict['raw_bytes'] = json_format.MessageToDict(
            account_pos, preserving_proto_field_name=True)
      elif msg_type == 'PARSED_NFT_TOKEN_INFO':
        nft_token_info = QueryNftTokenInfoResponseProto()
        nft_token_info.ParseFromString(raw_bytes)
        raw_msg_dict['raw_bytes'] = json_format.MessageToDict(
            nft_token_info, preserving_proto_field_name=True)
      elif msg_type == 'PARSED_FUNDING_FEE':
        funding_fee_bundle = FundingFeeBundle()
        funding_fee_bundle.ParseFromString(raw_bytes)
        raw_msg_dict['raw_bytes'] = json_format.MessageToDict(
            funding_fee_bundle, preserving_proto_field_name=True)
      elif msg_type == 'GRPC_MESSAGE':
        raw_msg_dict['raw_bytes'] = json.loads(raw_bytes)
      elif msg_type == 'PARSED_ORDER_FILL':
        fill_list = ProductFillBundle()
        fill_list.ParseFromString(raw_bytes)
        raw_msg_dict['raw_bytes'] = json_format.MessageToDict(
            fill_list, preserving_proto_field_name=True)
      elif msg_type == 'PARSED_ORDER_STATUS':
        order_list = ProductOrderBundle()
        order_list.ParseFromString(raw_bytes)
        raw_msg_dict['raw_bytes'] = json_format.MessageToDict(
            order_list, preserving_proto_field_name=True)
      elif msg_type == 'PARSED_SUBMIT_RESPONSE':
        submit_resp = SubmitOrderResponse()
        submit_resp.ParseFromString(raw_bytes)
        raw_msg_dict['raw_bytes'] = json_format.MessageToDict(
            submit_resp, preserving_proto_field_name=True)
      elif msg_type == 'PARSED_CANCEL_RESPONSE':
        cancel_resp = CancelOrderResponse()
        cancel_resp.ParseFromString(raw_bytes)
        raw_msg_dict['raw_bytes'] = json_format.MessageToDict(
            cancel_resp, preserving_proto_field_name=True)
      elif msg_type in ('SUBMIT_ORDER', 'CANCEL_ORDER', 'AMEND_ORDER', 'MULTI_CANCEL_ORDER'):
        order_execution = OrderExecutionProto()
        order_execution.ParseFromString(raw_bytes)
        raw_msg_dict['raw_bytes'] = json_format.MessageToDict(
            order_execution, preserving_proto_field_name=True)
      elif msg_type == 'PARSED_CONVERSION':
        conv_list = ConversionBundle()
        conv_list.ParseFromString(raw_bytes)
        raw_msg_dict['raw_bytes'] = json_format.MessageToDict(
            conv_list, preserving_proto_field_name=True)
      elif msg_type == 'PARSED_MULTI_SUBMIT_RESPONSE':
        resp = MultiSubmitResponse()
        resp.ParseFromString(raw_bytes)
        raw_msg_dict['raw_bytes'] = json_format.MessageToDict(
            resp, preserving_proto_field_name=True)
      elif msg_type == 'PARSED_MULTI_CANCEL_RESPONSE':
        resp = MultiCancelResponse()
        resp.ParseFromString(raw_bytes)
        raw_msg_dict['raw_bytes'] = json_format.MessageToDict(
            resp, preserving_proto_field_name=True)
      elif msg_type in (
          'REST_MESSAGE',
          'REST_ACCOUNT_BALANCE',
          'REST_ACCOUNT_POSITION',
          'REST_ALL_ORDERS',
          'REST_OPEN_ORDERS',
          'REST_FILL_ORDERS',
          'SUBMIT_RESPONSE',
          'CANCEL_RESPONSE',
          'LISTEN_KEY',
          'MULTI_CANCEL_RESPONSE'):
        pass
      elif msg_type == 'WS_RESET':
        error_code = ErrorCode()
        error_code.ParseFromString(raw_bytes)
        raw_msg_dict['raw_bytes'] = json_format.MessageToDict(
            error_code, preserving_proto_field_name=True)
      elif msg_type == 'WS_MESSAGE':
        try:
          raw_msg_dict['raw_bytes'] = json.loads(raw_bytes)
        except Exception:
          print('ws raw_byte: ', raw_bytes.decode())
          raw_msg_dict['raw_bytes'] = raw_bytes.decode()
      elif msg_type in ('ORDER_EXECUTION', 'SYSTEM'):
        pass
      else:
        print("unknown msg_type: ", msg_type)
        pass

      # deprecated!
      if raw_msg.HasField('http_response'):
        try:
          raw_msg_dict['http_response']['body'] = json.loads(res_body)
        except Exception:
          print('response body: ', res_body)

        if raw_msg.http_response.HasField('http_request'):
          try:
            raw_msg_dict['http_response']['http_request']['body'] = json.loads(req_body)
          except Exception:
            if req_body:
              print('request body: ', req_body)

      # replace above
      if raw_msg.HasField('http_context'):
        try:
          raw_msg_dict['http_context']['http_response']['body'] = json.loads(res_body)
        except Exception:
          print('response body: ', res_body)
          raw_msg_dict['http_context']['http_response']['body'] = res_body.decode()

        try:
          raw_msg_dict['http_context']['http_request']['body'] = json.loads(req_body)
        except Exception:
          if req_body:
            print('request body: ', req_body)
            raw_msg_dict['http_context']['http_request']['body'] = req_body.decode()
      pb.ClearField('private_exchange_message')

    data = json_format.MessageToDict(pb, preserving_proto_field_name=True)
    hr_ts = datetime.datetime.fromtimestamp(timestamp / 1e9).strftime('%Y%m%d %H:%M:%S.%f')

    data['source'] = source
    data['record_timestamp'] = timestamp
    data['record_timestamp_human_readable'] = hr_ts

    if pb.type == OrderGatewayLog.PRIVATE_EXCHANGE_MESSAGE:
      data['private_exchange_message'] = raw_msg_dict

    try:
      json.dump(data, self.file_obj)
    except ValueError:
      print(data)
      raise

    self.file_obj.write("\n")
    self.num += 1
    if timestamp > self.last_print_time + int(3600 * 1e9):
      print(self.num, hr_ts)
      self.last_print_time = timestamp


def launch_og_log(owner, market_type, exchange, trading_date, root_dir, machine, on_log):
  start_time = trading_date  # UTC 0 hour
  hours = 24

  run_from_og_log_archive(on_log,
                          market_type,
                          exchange,
                          start_time,
                          start_time + datetime.timedelta(hours=hours),
                          root_dir,
                          machine,
                          owner=owner)
  return 0


def main(_):
  owner = FLAGS.owner
  market_type = FLAGS.market_type
  exchange = FLAGS.exchange
  trading_date = FLAGS.trading_date
  root_dir = os.path.expanduser(FLAGS.root_dir)
  machine = FLAGS.machine
  assert market_type, '--market_type must be specified.'
  assert exchange, '--exchange must be specified.'
  assert trading_date, '--trading_date must be specified.'
  assert machine, '--machine must be specified.'
  print('Running for %s %s ...' % (trading_date, machine))
  sys.stdout.flush()

  filename = 'raw_og_log.%s.%s.%s.%s.json' % (
      FLAGS.owner, FLAGS.market_type, FLAGS.exchange, FLAGS.trading_date)

  pretty_filename = 'raw_og_log.pretty.%s.%s.%s.%s.json' % (
      FLAGS.owner, FLAGS.market_type, FLAGS.exchange, FLAGS.trading_date)

  onlog = OnOpLog(filename)
  trading_date = datetime.datetime.strptime(trading_date, '%Y%m%d')
  launch_og_log(owner, market_type, exchange, trading_date, root_dir, machine, onlog.on_log)
  onlog.close()

  out = open(pretty_filename, "wb")
  subprocess.run(["jq", "-s", '.', filename], stdout=out, check=False)
  out.close()


if __name__ == '__main__':
  flags.DEFINE_string('owner', None, 'Account.')
  flags.DEFINE_string('market_type', None, 'Market type.')
  flags.DEFINE_string('exchange', None, 'Exchange name.')
  flags.DEFINE_string('trading_date', None, 'Trading date in form of %Y%m%d.')
  flags.DEFINE_string('root_dir', '~/data/raw_og_log', 'root_dir.')
  flags.DEFINE_string('machine', None, 'Instance machine name.')
  app.run(main)
